//arch/riscv/kernel/proc.c
#include "defs.h"
#include "mm.h"
#include "proc.h"
#include "rand.h"
#include "printk.h"
#include "vm.h"
#include "elf.h"
#include "string.h"

extern uint64 uapp_start;
extern uint64 uapp_end;
extern uint64 swapper_pg_dir[512];


extern void __dummy();
extern void __switch_to(struct task_struct* prev, struct task_struct* next);

struct task_struct* idle;           // idle process
struct task_struct* current;        // 指向当前运行线程的 `task_struct`
struct task_struct* task[NR_TASKS]; // 线程数组，所有的线程都保存在此


void  memcpy (void* dst,void*  src,size_t count)
{
    char* tmp = dst;
    const char* s = src;
    while(count--){
        *tmp++ = *s++;
    }
}

static uint64_t load_program(struct task_struct* task) {

    uint64 ustart = (uint64)&uapp_start;
    uint64 uend = (uint64)&uapp_end;
    uint64 userSize = uend-ustart;

    Elf64_Ehdr* ehdr = (Elf64_Ehdr*)ustart;
    printk("ehdr:0x%016lx\n",ehdr);

    uint64_t phdr_start = (uint64_t)ehdr + ehdr->e_phoff;
    printk("phdr_start:0x%016lx\n",phdr_start);

    int phdr_cnt = ehdr->e_phnum;
    printk("phdr_cnt:%d\n",phdr_cnt);

    uint64* pgd_va = ((uint64)(task->pgd_satp) << 12) + (uint64)PA2VA_OFFSET;
    printk("pgd_satp: 0x%016lx\n",task->pgd_satp);
    printk("pgd_pa: 0x%016lx\n",(uint64)(task->pgd_satp) << 12);
    printk("pgd_va: 0x%016lx\n",pgd_va);
    for(int i = 0; i < 512; i++){
        pgd_va[i] = swapper_pg_dir[i];
    }

    Elf64_Phdr* phdr;
    for (int i = 0; i < phdr_cnt; i++) {
        phdr = (Elf64_Phdr*)(phdr_start + sizeof(Elf64_Phdr) * i);
        if(phdr->p_type == PT_LOAD){
            uint64 pg_num = ((phdr->p_vaddr&0xfff) + phdr->p_memsz - 1) / PGSIZE + 1;
            printk("pg_num:%d\n",pg_num);
            uint64 uapp_new = alloc_pages(pg_num);
            printk("uapp_new: 0x%016lx\n",uapp_new);
            printk("uapp_new_pa: 0x%016lx\n",VA2PA(uapp_new));
            printk("uapp_start:0x%016lx\n",&uapp_start);
            printk("p_offset = 0x%016lx\n",phdr->p_offset);
            uint64 load_addr = (uint64)(&uapp_start) + (uint64)phdr->p_offset;
            printk("load_addr:0x%016lx\n",load_addr);
            memcpy((uint64)((void*)(uapp_new) + (phdr->p_vaddr & 0xfff)), load_addr, (phdr->p_vaddr&0xfff) + phdr->p_memsz);
            printk("memcpy finished");
            printk("pgd_Va:0x%016lx\n",pgd_va);
            printk("p_vaddr:0x%016lx\n",phdr->p_vaddr);
            printk("uapp_pa:0x%016lx\n",uapp_new - (uint64)PA2VA_OFFSET);
            create_mapping((uint64*)(pgd_va), phdr->p_vaddr, uapp_new - (uint64)PA2VA_OFFSET, pg_num*PGSIZE, 0x8|phdr->p_flags);
            printk("map finished");
        }
        // following code has been written for you
        // set user stack
        uint64 user_stack = (uint64)alloc_page();
        printk("start mapping user stack\n");
        printk("user_stack:0x%016lx\n",user_stack);
        create_mapping(pgd_va, USER_END - PGSIZE, user_stack - (uint64)PA2VA_OFFSET, PGSIZE, 0xf);
        //...;
        // pc for the user program
        printk("sepc: 0x%016lx\n",(uint64)ehdr->e_entry);
        task->thread.sepc = (uint64)ehdr->e_entry;
        // sstatus bits set
        task->thread.sstatus &= ~(0x0000000000000100); //set spp = 0
        task->thread.sstatus |= 0x0000000000040020; //set spie and sum = 1
        printk("sstatus:0x%016lx\n",task->thread.sstatus);
        // user stack for user program
        task->thread.sscratch = USER_END;
    }
}

void task_init() {
    // 1. 调用 kalloc() 为 idle 分配一个物理页
    idle = (struct task_struct*)kalloc();
    // 2. 设置 state 为 TASK_RUNNING;
    idle -> state = TASK_RUNNING;
    // 3. 由于 idle 不参与调度 可以将其 counter / priority 设置为 0
    idle -> counter = 0;
    idle -> priority = 0;
    // 4. 设置 idle 的 pid 为 0
    idle -> pid = 0;
    idle ->thread.sscratch = 0;
    // 5. 将 current 和 task[0] 指向 idle
    current = idle;
    task[0] = idle;


    /* YOUR CODE HERE */
    for(int i = 1; i < NR_TASKS; i++){
        // 1. 参考 idle 的设置, 为 task[1] ~ task[NR_TASKS - 1] 进行初始化
        task[i] = (struct task_struct*)kalloc();
        // 2. 其中每个线程的 state 为 TASK_RUNNING, counter 为 0, priority 使用 rand() 来设置, pid 为该线程在线程数组中的下标。
        task[i]->state = TASK_RUNNING;
        task[i]->counter = 0;
        task[i]->priority = rand();
        task[i]->pid = i;

        // 3. 为 task[1] ~ task[NR_TASKS - 1] 设置 `thread_struct` 中的 `ra` 和 `sp`,
        // 4. 其中 `ra` 设置为 __dummy （见 4.3.2）的地址， `sp` 设置为 该线程申请的物理页的高地址
        task[i]->thread.ra = (uint64)__dummy;
        task[i]->thread.sp = (uint64)task[i] + PGSIZE;

        pagetable_t pgd_va = (pagetable_t)kalloc();
        printk("pgd_va:0x%016lx\n",pgd_va);

        uint64 pgd_pa = (uint64)pgd_va - (uint64)PA2VA_OFFSET;
        uint64 pgd_satp = (pgd_pa >> 12) | 0x8000000000000000;
        task[i]->pgd_satp = (pagetable_t)(pgd_satp);
        


        uint64 ustart = (uint64)&uapp_start;
        uint64 uend = (uint64)&uapp_end;

        load_program(task[i]);
        // printk("H\n");
        // U | X | W | R
        // 1 | 1 | 0 | 1
    }

    /* YOUR CODE HERE */

    printk("...proc_init done!\n");
}

void switch_to(struct task_struct* next) {
    /* YOUR CODE HERE */
    struct task_struct* prev = current;
    if(current != next){
        current = next;
        //printk("PREPARING to switch");
        __switch_to(prev, next);
    }
}


void do_timer(void) {
    /* 1. 如果当前线程是 idle 线程 直接进行调度 */
    if(current == idle){
        schedule();
    }
    /* 2. 如果当前线程不是 idle 对当前线程的运行剩余时间减 1 
          若剩余时间任然大于0 则直接返回 否则进行调度 */
    else{
        current->counter--;
        if(current->counter > 0){
            return;
        }
        schedule();
    }
    /* YOUR CODE HERE */
}

#ifdef SJF

void schedule(void){
    int min_time_index = 1;
    int min_time = MAX_COUNTER;
    int flag = 0;
    for(int i = 1; i < NR_TASKS; i++){
        if(task[i]->state == TASK_RUNNING){
            if(task[i]->counter > 0){
                flag = 1;
                /*Find the task with the minimum counter*/
                if(task[i]->counter < min_time){
                    min_time_index = i;
                    min_time = task[i]->counter;
                }
            }
        }
    }
    if(!flag){
        printk("\n");
        /* reset the remaining time of the task */
        for(int i = 1; i < NR_TASKS; i++){
            task[i]->counter = rand();
            printk("SET [PID = %d COUNTER = %d]\n",i,task[i]->counter);
        }
        schedule();
    } else {
        /* switch to a new task if the current one has been done*/
        printk("\nswitch to [PID = %d COUNTER = %d]\n",min_time_index,min_time);
        switch_to(task[min_time_index]);
        //printk("\nreturn from switch_to")
    }

}

#endif

#ifdef PRIORITY

void schedule(void){
    /*Step 1: Find the max counter and the matching task_index*/
    int maxCounter = 0;
    int maxCounterIndex = 0;
    for(int i = 1; i < NR_TASKS; i++){
        if(task[i]->state == TASK_RUNNING && task[i]->counter > maxCounter){
            maxCounter = task[i]->counter;
            maxCounterIndex = i;
        }
    }
    /*Step 2: Find the task to be executed (that should be the one with the longest counter)*/
    if(maxCounter){
        printk("\nswitch to [PID = %d PRIORITY = %d COUNTER = %d]\n",maxCounterIndex,task[maxCounterIndex]->priority,task[maxCounterIndex]->counter);
        switch_to(task[maxCounterIndex]);
        return;
    }
    /*Step 3: If max counter is 0, randomize task again, and reschedule*/
    {
        /*reset the task's counter considering its priority*/
        printk("\n");
        for(int i = 1; i < NR_TASKS; i++){
            task[i]->counter = (task[i]->counter >> 1) + (task[i]->priority);
            /*Display the changes*/
            printk("SET [PID = %d PRIORITY = %d COUNTER = %d]\n",i,task[i]->priority,task[i]->counter);
        }
        /*reschedule*/
        schedule();
    }
}

#endif


void dummy() {
    uint64 MOD = 1000000007;
    uint64 auto_inc_local_var = 0;
    int last_counter = -1;
    while(1) {
        if (last_counter == -1 || current->counter != last_counter) {
            last_counter = current->counter;
            auto_inc_local_var = (auto_inc_local_var + 1) % MOD;
            printk("[PID = %d] is running. auto_inc_local_var = %d\n", current->pid, auto_inc_local_var);
        }
    }
}

extern void __dummy(){
    csr_write(sepc, current->thread.sepc);
    asm volatile("csrrw sp, sscratch, sp");
    asm volatile("sret");
}